/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia;



import com.google.auto.service.AutoService;

import com.bff.gaia.unified.runners.core.KeyedWorkItem;

import com.bff.gaia.unified.runners.core.SplittableParDoViaKeyedWorkItems;

import com.bff.gaia.unified.runners.core.SystemReduceFn;

import com.bff.gaia.unified.runners.core.construction.PTransformTranslation;

import com.bff.gaia.unified.runners.core.construction.ParDoTranslation;

import com.bff.gaia.unified.runners.core.construction.ReadTranslation;

import com.bff.gaia.unified.runners.core.construction.SplittableParDo;

import com.bff.gaia.unified.runners.core.construction.TransformPayloadTranslatorRegistrar;

import com.bff.gaia.unified.runners.core.construction.UnboundedReadFromBoundedSource.BoundedToUnboundedSourceAdapter;

import com.bff.gaia.unified.runners.gaia.translation.functions.GaiaAssignWindows;

import com.bff.gaia.unified.runners.gaia.translation.types.CoderTypeInformation;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.*;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.DoFnOperator;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.KvToByteBufferKeySelector;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.SingletonKeyedWorkItem;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.SingletonKeyedWorkItemCoder;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.SplittableDoFnOperator;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.WindowDoFnOperator;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.WorkItemKeySelector;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.io.DedupingOperator;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.io.TestStreamSource;

import com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.io.UnboundedSourceWrapper;

import com.bff.gaia.unified.sdk.coders.ByteArrayCoder;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.CoderException;

import com.bff.gaia.unified.sdk.coders.KvCoder;

import com.bff.gaia.unified.sdk.coders.VoidCoder;

import com.bff.gaia.unified.sdk.io.BoundedSource;

import com.bff.gaia.unified.sdk.io.UnboundedSource;

import com.bff.gaia.unified.sdk.runners.AppliedPTransform;

import com.bff.gaia.unified.sdk.testing.TestStream;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.CombineFnBase.GlobalCombineFn;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.DoFnSchemaInformation;

import com.bff.gaia.unified.sdk.transforms.PTransform;

import com.bff.gaia.unified.sdk.transforms.SerializableFunction;

import com.bff.gaia.unified.sdk.transforms.join.RawUnionValue;

import com.bff.gaia.unified.sdk.transforms.join.UnionCoder;

import com.bff.gaia.unified.sdk.transforms.reflect.DoFnSignature;

import com.bff.gaia.unified.sdk.transforms.reflect.DoFnSignatures;

import com.bff.gaia.unified.sdk.transforms.windowing.BoundedWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.GlobalWindow;

import com.bff.gaia.unified.sdk.transforms.windowing.WindowFn;

import com.bff.gaia.unified.sdk.util.AppliedCombineFn;

import com.bff.gaia.unified.sdk.util.CoderUtils;

import com.bff.gaia.unified.sdk.util.WindowedValue;

import com.bff.gaia.unified.vendor.guava.com.google.common.annotations.VisibleForTesting;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.ImmutableMap;

import com.bff.gaia.unified.vendor.guava.com.google.common.collect.Maps;

import com.bff.gaia.api.common.functions.FlatMapFunction;

import com.bff.gaia.api.common.functions.MapFunction;

import com.bff.gaia.api.common.functions.RichFlatMapFunction;

import com.bff.gaia.api.common.functions.StoppableFunction;

import com.bff.gaia.api.common.typeinfo.TypeInformation;

import com.bff.gaia.api.java.functions.KeySelector;

import com.bff.gaia.api.java.tuple.Tuple2;

import com.bff.gaia.api.java.typeutils.GenericTypeInfo;

import com.bff.gaia.api.java.typeutils.ResultTypeQueryable;

import com.bff.gaia.unified.sdk.values.KV;

import com.bff.gaia.unified.sdk.values.PBegin;

import com.bff.gaia.unified.sdk.values.PCollection;

import com.bff.gaia.unified.sdk.values.PCollectionTuple;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.PValue;

import com.bff.gaia.unified.sdk.values.TupleTag;

import com.bff.gaia.unified.sdk.values.TupleTagList;

import com.bff.gaia.unified.sdk.values.ValueWithRecordId;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;

import com.bff.gaia.configuration.Configuration;

import com.bff.gaia.runtime.state.CheckpointListener;

import com.bff.gaia.runtime.state.FunctionInitializationContext;

import com.bff.gaia.runtime.state.FunctionSnapshotContext;

import com.bff.gaia.streaming.api.checkpoint.CheckpointedFunction;

import com.bff.gaia.streaming.api.datastream.DataStream;

import com.bff.gaia.streaming.api.datastream.DataStreamSource;

import com.bff.gaia.streaming.api.datastream.KeyedStream;

import com.bff.gaia.streaming.api.datastream.SingleOutputStreamOperator;

import com.bff.gaia.streaming.api.functions.source.RichParallelSourceFunction;

import com.bff.gaia.streaming.api.operators.OneInputStreamOperator;

import com.bff.gaia.streaming.api.operators.TwoInputStreamOperator;

import com.bff.gaia.streaming.api.transformations.TwoInputTransformation;

import com.bff.gaia.streaming.api.watermark.Watermark;

import com.bff.gaia.streaming.runtime.tasks.ProcessingTimeCallback;

import com.bff.gaia.util.Collector;

import com.bff.gaia.util.OutputTag;



import javax.annotation.Nullable;

import java.io.IOException;

import java.nio.ByteBuffer;

import java.util.*;



import static java.lang.String.format;

import static com.bff.gaia.unified.runners.core.construction.SplittableParDo.SPLITTABLE_PROCESS_URN;



/**

 * This class contains all the mappings between Unified and Gaia <b>streaming</b> transformations. The

 * {@link GaiaStreamingPipelineTranslator} traverses the Unified job and comes here to translate the

 * encountered Unified transformations into Gaia one, based on the mapping available in this class.

 */

class GaiaStreamingTransformTranslators {



  // --------------------------------------------------------------------------------------------

  //  Transform Translator Registry

  // --------------------------------------------------------------------------------------------



  /** A map from a Transform URN to the translator. */

  @SuppressWarnings("rawtypes")

  private static final Map<String, GaiaStreamingPipelineTranslator.StreamTransformTranslator>

      TRANSLATORS = new HashMap<>();



  // here you can find all the available translators.

  static {

    TRANSLATORS.put(PTransformTranslation.READ_TRANSFORM_URN, new ReadSourceTranslator());



    TRANSLATORS.put(PTransformTranslation.PAR_DO_TRANSFORM_URN, new ParDoStreamingTranslator());

    TRANSLATORS.put(SPLITTABLE_PROCESS_URN, new SplittableProcessElementsStreamingTranslator());

    TRANSLATORS.put(SplittableParDo.SPLITTABLE_GBKIKWI_URN, new GBKIntoKeyedWorkItemsTranslator());



    TRANSLATORS.put(

        PTransformTranslation.ASSIGN_WINDOWS_TRANSFORM_URN, new WindowAssignTranslator());

    TRANSLATORS.put(

        PTransformTranslation.FLATTEN_TRANSFORM_URN, new FlattenPCollectionTranslator());

    TRANSLATORS.put(

        CreateStreamingGaiaView.CREATE_STREAMING_GAIA_VIEW_URN,

        new CreateViewStreamingTranslator());



    TRANSLATORS.put(PTransformTranslation.RESHUFFLE_URN, new ReshuffleTranslatorStreaming());

    TRANSLATORS.put(PTransformTranslation.GROUP_BY_KEY_TRANSFORM_URN, new GroupByKeyTranslator());

    TRANSLATORS.put(

        PTransformTranslation.COMBINE_PER_KEY_TRANSFORM_URN, new CombinePerKeyTranslator());



    TRANSLATORS.put(PTransformTranslation.TEST_STREAM_TRANSFORM_URN, new TestStreamTranslator());

  }



  public static GaiaStreamingPipelineTranslator.StreamTransformTranslator<?> getTranslator(

      PTransform<?, ?> transform) {

    @Nullable String urn = PTransformTranslation.urnForTransformOrNull(transform);

    return urn == null ? null : TRANSLATORS.get(urn);

  }



  @SuppressWarnings("unchecked")

  private static String getCurrentTransformName(GaiaStreamingTranslationContext context) {

    return context.getCurrentTransform().getFullName();

  }



  // --------------------------------------------------------------------------------------------

  //  Transformation Implementations

  // --------------------------------------------------------------------------------------------



  private static class UnboundedReadSourceTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PBegin, PCollection<T>>> {



    @Override

    public void translateNode(

        PTransform<PBegin, PCollection<T>> transform, GaiaStreamingTranslationContext context) {

      PCollection<T> output = context.getOutput(transform);



      DataStream<WindowedValue<T>> source;

      DataStream<WindowedValue<ValueWithRecordId<T>>> nonDedupSource;

      TypeInformation<WindowedValue<T>> outputTypeInfo =

          context.getTypeInfo(context.getOutput(transform));



      Coder<T> coder = context.getOutput(transform).getCoder();



      TypeInformation<WindowedValue<ValueWithRecordId<T>>> withIdTypeInfo =

          new CoderTypeInformation<>(

              WindowedValue.getFullCoder(

                  ValueWithRecordId.ValueWithRecordIdCoder.of(coder),

                  output.getWindowingStrategy().getWindowFn().windowCoder()));



      UnboundedSource<T, ?> rawSource;

      try {

        rawSource =

            ReadTranslation.unboundedSourceFromTransform(

                (AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>>)

                    context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      String fullName = getCurrentTransformName(context);

      try {

        int parallelism =

            context.getExecutionEnvironment().getMaxParallelism() > 0

                ? context.getExecutionEnvironment().getMaxParallelism()

                : context.getExecutionEnvironment().getParallelism();

        UnboundedSourceWrapper<T, ?> sourceWrapper =

            new UnboundedSourceWrapper<>(

                fullName, context.getPipelineOptions(), rawSource, parallelism);

        nonDedupSource =

            context

                .getExecutionEnvironment()

                .addSource(sourceWrapper)

                .name(fullName)

                .uid(fullName)

                .returns(withIdTypeInfo);



        if (rawSource.requiresDeduping()) {

          source =

              nonDedupSource

                  .keyBy(new ValueWithRecordIdKeySelector<>())

                  .transform("deduping", outputTypeInfo, new DedupingOperator<>())

                  .uid(format("%s/__deduplicated__", fullName));

        } else {

          source = nonDedupSource.flatMap(new StripIdsMap<>()).returns(outputTypeInfo);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error while translating UnboundedSource: " + rawSource, e);

      }



      context.setOutputDataStream(output, source);

    }

  }



  static class ValueWithRecordIdKeySelector<T>

      implements KeySelector<WindowedValue<ValueWithRecordId<T>>, ByteBuffer>,

          ResultTypeQueryable<ByteBuffer> {



    @Override

    public ByteBuffer getKey(WindowedValue<ValueWithRecordId<T>> value) throws Exception {

      return ByteBuffer.wrap(value.getValue().getId());

    }



    @Override

    public TypeInformation<ByteBuffer> getProducedType() {

      return new GenericTypeInfo<>(ByteBuffer.class);

    }

  }



  public static class StripIdsMap<T>

      implements FlatMapFunction<WindowedValue<ValueWithRecordId<T>>, WindowedValue<T>> {



    @Override

    public void flatMap(

        WindowedValue<ValueWithRecordId<T>> value, Collector<WindowedValue<T>> collector)

        throws Exception {

      collector.collect(value.withValue(value.getValue().getValue()));

    }

  }



  private static class ReadSourceTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PBegin, PCollection<T>>> {



    private final BoundedReadSourceTranslator<T> boundedTranslator =

        new BoundedReadSourceTranslator<>();

    private final UnboundedReadSourceTranslator<T> unboundedTranslator =

        new UnboundedReadSourceTranslator<>();



    @Override

    void translateNode(

        PTransform<PBegin, PCollection<T>> transform, GaiaStreamingTranslationContext context) {

      if (context.getOutput(transform).isBounded().equals(PCollection.IsBounded.BOUNDED)) {

        boundedTranslator.translateNode(transform, context);

      } else {

        unboundedTranslator.translateNode(transform, context);

      }

    }

  }



  private static class BoundedReadSourceTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PBegin, PCollection<T>>> {



    @Override

    public void translateNode(

        PTransform<PBegin, PCollection<T>> transform, GaiaStreamingTranslationContext context) {

      PCollection<T> output = context.getOutput(transform);



      TypeInformation<WindowedValue<T>> outputTypeInfo =

          context.getTypeInfo(context.getOutput(transform));



      BoundedSource<T> rawSource;

      try {

        rawSource =

            ReadTranslation.boundedSourceFromTransform(

                (AppliedPTransform<PBegin, PCollection<T>, PTransform<PBegin, PCollection<T>>>)

                    context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      String fullName = getCurrentTransformName(context);

      UnboundedSource<T, ?> adaptedRawSource = new BoundedToUnboundedSourceAdapter<>(rawSource);

      DataStream<WindowedValue<T>> source;

      try {

        int parallelism =

            context.getExecutionEnvironment().getMaxParallelism() > 0

                ? context.getExecutionEnvironment().getMaxParallelism()

                : context.getExecutionEnvironment().getParallelism();

        UnboundedSourceWrapperNoValueWithRecordId<T, ?> sourceWrapper =

            new UnboundedSourceWrapperNoValueWithRecordId<>(

                new UnboundedSourceWrapper<>(

                    fullName, context.getPipelineOptions(), adaptedRawSource, parallelism));

        source =

            context

                .getExecutionEnvironment()

                .addSource(sourceWrapper)

                .name(fullName)

                .uid(fullName)

                .returns(outputTypeInfo);

      } catch (Exception e) {

        throw new RuntimeException("Error while translating BoundedSource: " + rawSource, e);

      }

      context.setOutputDataStream(output, source);

    }

  }



  /** Wraps each element in a {@link RawUnionValue} with the given tag id. */

  public static class ToRawUnion<T> implements MapFunction<T, RawUnionValue> {

    private final int intTag;



    public ToRawUnion(int intTag) {

      this.intTag = intTag;

    }



    @Override

    public RawUnionValue map(T o) throws Exception {

      return new RawUnionValue(intTag, o);

    }

  }



  private static Tuple2<Map<Integer, PCollectionView<?>>, DataStream<RawUnionValue>>

      transformSideInputs(

          Collection<PCollectionView<?>> sideInputs, GaiaStreamingTranslationContext context) {



    // collect all side inputs

    Map<TupleTag<?>, Integer> tagToIntMapping = new HashMap<>();

    Map<Integer, PCollectionView<?>> intToViewMapping = new HashMap<>();

    int count = 0;

    for (PCollectionView<?> sideInput : sideInputs) {

      TupleTag<?> tag = sideInput.getTagInternal();

      intToViewMapping.put(count, sideInput);

      tagToIntMapping.put(tag, count);

      count++;

    }



    List<Coder<?>> inputCoders = new ArrayList<>();

    for (PCollectionView<?> sideInput : sideInputs) {

      DataStream<Object> sideInputStream = context.getInputDataStream(sideInput);

      TypeInformation<Object> tpe = sideInputStream.getType();

      if (!(tpe instanceof CoderTypeInformation)) {

        throw new IllegalStateException("Input Stream TypeInformation is no CoderTypeInformation.");

      }



      Coder<?> coder = ((CoderTypeInformation) tpe).getCoder();

      inputCoders.add(coder);

    }



    UnionCoder unionCoder = UnionCoder.of(inputCoders);



    CoderTypeInformation<RawUnionValue> unionTypeInformation =

        new CoderTypeInformation<>(unionCoder);



    // transform each side input to RawUnionValue and union them

    DataStream<RawUnionValue> sideInputUnion = null;



    for (PCollectionView<?> sideInput : sideInputs) {

      TupleTag<?> tag = sideInput.getTagInternal();

      final int intTag = tagToIntMapping.get(tag);

      DataStream<Object> sideInputStream = context.getInputDataStream(sideInput);

      DataStream<RawUnionValue> unionValueStream =

          sideInputStream.map(new ToRawUnion<>(intTag)).returns(unionTypeInformation);



      if (sideInputUnion == null) {

        sideInputUnion = unionValueStream;

      } else {

        sideInputUnion = sideInputUnion.union(unionValueStream);

      }

    }



    if (sideInputUnion == null) {

      throw new IllegalStateException("No unioned side inputs, this indicates a bug.");

    }



    return new Tuple2<>(intToViewMapping, sideInputUnion);

  }



  /**

   * Helper for translating {@code ParDo.MultiOutput} and {@link

   * SplittableParDoViaKeyedWorkItems.ProcessElements}.

   */

  static class ParDoTranslationHelper {



    interface DoFnOperatorFactory<InputT, OutputT> {

      DoFnOperator<InputT, OutputT> createDoFnOperator(

		  DoFn<InputT, OutputT> doFn,

		  String stepName,

		  List<PCollectionView<?>> sideInputs,

		  TupleTag<OutputT> mainOutputTag,

		  List<TupleTag<?>> additionalOutputTags,

		  GaiaStreamingTranslationContext context,

		  WindowingStrategy<?, ?> windowingStrategy,

		  Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags,

		  Map<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders,

		  Map<TupleTag<?>, Integer> tagsToIds,

		  Coder<WindowedValue<InputT>> windowedInputCoder,

		  Coder<InputT> inputCoder,

		  Map<TupleTag<?>, Coder<?>> outputCoders,

		  Coder keyCoder,

		  KeySelector<WindowedValue<InputT>, ?> keySelector,

		  Map<Integer, PCollectionView<?>> transformedSideInputs,

		  DoFnSchemaInformation doFnSchemaInformation);

    }



    static <InputT, OutputT> void translateParDo(

        String transformName,

        DoFn<InputT, OutputT> doFn,

        PCollection<InputT> input,

        List<PCollectionView<?>> sideInputs,

        Map<TupleTag<?>, PValue> outputs,

        TupleTag<OutputT> mainOutputTag,

        List<TupleTag<?>> additionalOutputTags,

        DoFnSchemaInformation doFnSchemaInformation,

        GaiaStreamingTranslationContext context,

        DoFnOperatorFactory<InputT, OutputT> doFnOperatorFactory) {



      // we assume that the transformation does not change the windowing strategy.

      WindowingStrategy<?, ?> windowingStrategy = input.getWindowingStrategy();



      Map<TupleTag<?>, OutputTag<WindowedValue<?>>> tagsToOutputTags = Maps.newHashMap();

      Map<TupleTag<?>, Coder<WindowedValue<?>>> tagsToCoders = Maps.newHashMap();



      // We associate output tags with ids, the Integer is easier to serialize than TupleTag.

      // The return map of AppliedPTransform.getOutputs() is an ImmutableMap, its implementation is

      // RegularImmutableMap, its entrySet order is the same with the order of insertion.

      // So we can use the original AppliedPTransform.getOutputs() to produce deterministic ids.

      Map<TupleTag<?>, Integer> tagsToIds = Maps.newHashMap();

      int idCount = 0;

      tagsToIds.put(mainOutputTag, idCount++);

      for (Map.Entry<TupleTag<?>, PValue> entry : outputs.entrySet()) {

        if (!tagsToOutputTags.containsKey(entry.getKey())) {

          tagsToOutputTags.put(

              entry.getKey(),

              new OutputTag<WindowedValue<?>>(

                  entry.getKey().getId(),

                  (TypeInformation) context.getTypeInfo((PCollection<?>) entry.getValue())));

          tagsToCoders.put(

              entry.getKey(),

              (Coder) context.getWindowedInputCoder((PCollection<OutputT>) entry.getValue()));

          tagsToIds.put(entry.getKey(), idCount++);

        }

      }



      SingleOutputStreamOperator<WindowedValue<OutputT>> outputStream;



      Coder<WindowedValue<InputT>> windowedInputCoder = context.getWindowedInputCoder(input);

      Coder<InputT> inputCoder = context.getInputCoder(input);

      Map<TupleTag<?>, Coder<?>> outputCoders = context.getOutputCoders();



      DataStream<WindowedValue<InputT>> inputDataStream = context.getInputDataStream(input);



      Coder keyCoder = null;

      KeySelector<WindowedValue<InputT>, ?> keySelector = null;

      boolean stateful = false;

      DoFnSignature signature = DoFnSignatures.getSignature(doFn.getClass());

      if (signature.stateDeclarations().size() > 0 || signature.timerDeclarations().size() > 0) {

        // Based on the fact that the signature is stateful, DoFnSignatures ensures

        // that it is also keyed

        keyCoder = ((KvCoder) input.getCoder()).getKeyCoder();

        keySelector = new KvToByteBufferKeySelector(keyCoder);

        inputDataStream = inputDataStream.keyBy(keySelector);

        stateful = true;

      } else if (doFn instanceof SplittableParDoViaKeyedWorkItems.ProcessFn) {

        // we know that it is keyed on byte[]

        keyCoder = ByteArrayCoder.of();

        keySelector = new WorkItemKeySelector<>(keyCoder);

        stateful = true;

      }



      CoderTypeInformation<WindowedValue<OutputT>> outputTypeInformation =

          new CoderTypeInformation<>(

              context.getWindowedInputCoder((PCollection<OutputT>) outputs.get(mainOutputTag)));



      if (sideInputs.isEmpty()) {

        DoFnOperator<InputT, OutputT> doFnOperator =

            doFnOperatorFactory.createDoFnOperator(

                doFn,

                getCurrentTransformName(context),

                sideInputs,

                mainOutputTag,

                additionalOutputTags,

                context,

                windowingStrategy,

                tagsToOutputTags,

                tagsToCoders,

                tagsToIds,

                windowedInputCoder,

                inputCoder,

                outputCoders,

                keyCoder,

                keySelector,

                new HashMap<>() /* side-input mapping */,

                doFnSchemaInformation);



        outputStream =

            inputDataStream.transform(transformName, outputTypeInformation, doFnOperator);



      } else {

        Tuple2<Map<Integer, PCollectionView<?>>, DataStream<RawUnionValue>> transformedSideInputs =

            transformSideInputs(sideInputs, context);



        DoFnOperator<InputT, OutputT> doFnOperator =

            doFnOperatorFactory.createDoFnOperator(

                doFn,

                getCurrentTransformName(context),

                sideInputs,

                mainOutputTag,

                additionalOutputTags,

                context,

                windowingStrategy,

                tagsToOutputTags,

                tagsToCoders,

                tagsToIds,

                windowedInputCoder,

                inputCoder,

                outputCoders,

                keyCoder,

                keySelector,

                transformedSideInputs.f0,

                doFnSchemaInformation);



        if (stateful) {

          // we have to manually construct the two-input transform because we're not

          // allowed to have only one input keyed, normally.

          KeyedStream keyedStream = (KeyedStream<?, InputT>) inputDataStream;

          TwoInputTransformation<

                  WindowedValue<KV<?, InputT>>, RawUnionValue, WindowedValue<OutputT>>

              rawGaiaTransform =

                  new TwoInputTransformation(

                      keyedStream.getTransformation(),

                      transformedSideInputs.f1.broadcast().getTransformation(),

                      transformName,

                      doFnOperator,

                      outputTypeInformation,

                      keyedStream.getParallelism());



          rawGaiaTransform.setStateKeyType(keyedStream.getKeyType());

          rawGaiaTransform.setStateKeySelectors(keyedStream.getKeySelector(), null);



          outputStream =

              new SingleOutputStreamOperator(

                  keyedStream.getExecutionEnvironment(),

                  rawGaiaTransform) {}; // we have to cheat around the ctor being protected



          keyedStream.getExecutionEnvironment().addOperator(rawGaiaTransform);



        } else {

          outputStream =

              inputDataStream

                  .connect(transformedSideInputs.f1.broadcast())

                  .transform(transformName, outputTypeInformation, doFnOperator);

        }

      }



      outputStream.uid(transformName);

      context.setOutputDataStream(outputs.get(mainOutputTag), outputStream);



      for (Map.Entry<TupleTag<?>, PValue> entry : outputs.entrySet()) {

        if (!entry.getKey().equals(mainOutputTag)) {

          context.setOutputDataStream(

              entry.getValue(), outputStream.getSideOutput(tagsToOutputTags.get(entry.getKey())));

        }

      }

    }

  }



  private static class ParDoStreamingTranslator<InputT, OutputT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<InputT>, PCollectionTuple>> {



    @Override

    public void translateNode(

        PTransform<PCollection<InputT>, PCollectionTuple> transform,

        GaiaStreamingTranslationContext context) {



      DoFn<InputT, OutputT> doFn;

      try {

        doFn = (DoFn<InputT, OutputT>) ParDoTranslation.getDoFn(context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      TupleTag<OutputT> mainOutputTag;

      try {

        mainOutputTag =

            (TupleTag<OutputT>) ParDoTranslation.getMainOutputTag(context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      List<PCollectionView<?>> sideInputs;

      try {

        sideInputs = ParDoTranslation.getSideInputs(context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      TupleTagList additionalOutputTags;

      try {

        additionalOutputTags =

            ParDoTranslation.getAdditionalOutputTags(context.getCurrentTransform());

      } catch (IOException e) {

        throw new RuntimeException(e);

      }



      DoFnSchemaInformation doFnSchemaInformation;

      doFnSchemaInformation = ParDoTranslation.getSchemaInformation(context.getCurrentTransform());



      ParDoTranslationHelper.translateParDo(

          getCurrentTransformName(context),

          doFn,

          context.getInput(transform),

          sideInputs,

          context.getOutputs(transform),

          mainOutputTag,

          additionalOutputTags.getAll(),

          doFnSchemaInformation,

          context,

          (doFn1,

              stepName,

              sideInputs1,

              mainOutputTag1,

              additionalOutputTags1,

              context1,

              windowingStrategy,

              tagsToOutputTags,

              tagsToCoders,

              tagsToIds,

              windowedInputCoder,

              inputCoder,

              outputCoders1,

              keyCoder,

              keySelector,

              transformedSideInputs,

              doFnSchemaInformation1) ->

              new DoFnOperator<>(

                  doFn1,

                  stepName,

                  windowedInputCoder,

                  inputCoder,

                  outputCoders1,

                  mainOutputTag1,

                  additionalOutputTags1,

                  new DoFnOperator.MultiOutputOutputManagerFactory<>(

                      mainOutputTag1, tagsToOutputTags, tagsToCoders, tagsToIds),

                  windowingStrategy,

                  transformedSideInputs,

                  sideInputs1,

                  context1.getPipelineOptions(),

                  keyCoder,

                  keySelector,

                  doFnSchemaInformation1));

    }

  }



  private static class SplittableProcessElementsStreamingTranslator<

          InputT, OutputT, RestrictionT, PositionT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          SplittableParDoViaKeyedWorkItems.ProcessElements<

              InputT, OutputT, RestrictionT, PositionT>> {



    @Override

    public void translateNode(

        SplittableParDoViaKeyedWorkItems.ProcessElements<InputT, OutputT, RestrictionT, PositionT>

            transform,

        GaiaStreamingTranslationContext context) {



      ParDoTranslationHelper.translateParDo(

          getCurrentTransformName(context),

          transform.newProcessFn(transform.getFn()),

          context.getInput(transform),

          transform.getSideInputs(),

          context.getOutputs(transform),

          transform.getMainOutputTag(),

          transform.getAdditionalOutputTags().getAll(),

          DoFnSchemaInformation.create(),

          context,

          (doFn,

              stepName,

              sideInputs,

              mainOutputTag,

              additionalOutputTags,

              context1,

              windowingStrategy,

              tagsToOutputTags,

              tagsToCoders,

              tagsToIds,

              windowedInputCoder,

              inputCoder,

              outputCoders1,

              keyCoder,

              keySelector,

              transformedSideInputs,

              doFnSchemaInformation) ->

              new SplittableDoFnOperator<>(

                  doFn,

                  stepName,

                  windowedInputCoder,

                  inputCoder,

                  outputCoders1,

                  mainOutputTag,

                  additionalOutputTags,

                  new DoFnOperator.MultiOutputOutputManagerFactory<>(

                      mainOutputTag, tagsToOutputTags, tagsToCoders, tagsToIds),

                  windowingStrategy,

                  transformedSideInputs,

                  sideInputs,

                  context1.getPipelineOptions(),

                  keyCoder,

                  keySelector));

    }

  }



  private static class CreateViewStreamingTranslator<ElemT, ViewT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          CreateStreamingGaiaView.CreateGaiaPCollectionView<ElemT, ViewT>> {



    @Override

    public void translateNode(

        CreateStreamingGaiaView.CreateGaiaPCollectionView<ElemT, ViewT> transform,

        GaiaStreamingTranslationContext context) {

      // just forward

      DataStream<WindowedValue<List<ElemT>>> inputDataSet =

          context.getInputDataStream(context.getInput(transform));



      PCollectionView<ViewT> view = transform.getView();



      context.setOutputDataStream(view, inputDataSet);

    }

  }



  private static class WindowAssignTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<T>, PCollection<T>>> {



    @Override

    public void translateNode(

        PTransform<PCollection<T>, PCollection<T>> transform,

        GaiaStreamingTranslationContext context) {



      @SuppressWarnings("unchecked")

      WindowingStrategy<T, BoundedWindow> windowingStrategy =

          (WindowingStrategy<T, BoundedWindow>) context.getOutput(transform).getWindowingStrategy();



      TypeInformation<WindowedValue<T>> typeInfo =

          context.getTypeInfo(context.getOutput(transform));



      DataStream<WindowedValue<T>> inputDataStream =

          context.getInputDataStream(context.getInput(transform));



      WindowFn<T, ? extends BoundedWindow> windowFn = windowingStrategy.getWindowFn();



      GaiaAssignWindows<T, ? extends BoundedWindow> assignWindowsFunction =

          new GaiaAssignWindows<>(windowFn);



      String fullName = context.getOutput(transform).getName();

      SingleOutputStreamOperator<WindowedValue<T>> outputDataStream =

          inputDataStream

              .flatMap(assignWindowsFunction)

              .name(fullName)

              .uid(fullName)

              .returns(typeInfo);



      context.setOutputDataStream(context.getOutput(transform), outputDataStream);

    }

  }



  private static class ReshuffleTranslatorStreaming<K, InputT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, InputT>>>> {



    @Override

    public void translateNode(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, InputT>>> transform,

        GaiaStreamingTranslationContext context) {



      DataStream<WindowedValue<KV<K, InputT>>> inputDataSet =

          context.getInputDataStream(context.getInput(transform));



      context.setOutputDataStream(context.getOutput(transform), inputDataSet.rebalance());

    }

  }



  private static class GroupByKeyTranslator<K, InputT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, Iterable<InputT>>>>> {



    @Override

    public void translateNode(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, Iterable<InputT>>>> transform,

        GaiaStreamingTranslationContext context) {



      PCollection<KV<K, InputT>> input = context.getInput(transform);



      @SuppressWarnings("unchecked")

      WindowingStrategy<?, BoundedWindow> windowingStrategy =

          (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();



      KvCoder<K, InputT> inputKvCoder = (KvCoder<K, InputT>) input.getCoder();



      SingletonKeyedWorkItemCoder<K, InputT> workItemCoder =

          SingletonKeyedWorkItemCoder.of(

              inputKvCoder.getKeyCoder(),

              inputKvCoder.getValueCoder(),

              input.getWindowingStrategy().getWindowFn().windowCoder());



      DataStream<WindowedValue<KV<K, InputT>>> inputDataStream = context.getInputDataStream(input);



      WindowedValue.FullWindowedValueCoder<SingletonKeyedWorkItem<K, InputT>>

          windowedWorkItemCoder =

              WindowedValue.getFullCoder(

                  workItemCoder, input.getWindowingStrategy().getWindowFn().windowCoder());



      CoderTypeInformation<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemTypeInfo =

          new CoderTypeInformation<>(windowedWorkItemCoder);



      DataStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemStream =

          inputDataStream

              .flatMap(new ToKeyedWorkItem<>())

              .returns(workItemTypeInfo)

              .name("ToKeyedWorkItem");



      WorkItemKeySelector keySelector = new WorkItemKeySelector<>(inputKvCoder.getKeyCoder());



      KeyedStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>, ByteBuffer>

          keyedWorkItemStream =

              workItemStream.keyBy(new WorkItemKeySelector<>(inputKvCoder.getKeyCoder()));



      SystemReduceFn<K, InputT, Iterable<InputT>, Iterable<InputT>, BoundedWindow> reduceFn =

          SystemReduceFn.buffering(inputKvCoder.getValueCoder());



      Coder<WindowedValue<KV<K, Iterable<InputT>>>> outputCoder =

          context.getWindowedInputCoder(context.getOutput(transform));

      TypeInformation<WindowedValue<KV<K, Iterable<InputT>>>> outputTypeInfo =

          context.getTypeInfo(context.getOutput(transform));



      TupleTag<KV<K, Iterable<InputT>>> mainTag = new TupleTag<>("main output");



      String fullName = getCurrentTransformName(context);

      WindowDoFnOperator<K, InputT, Iterable<InputT>> doFnOperator =

          new WindowDoFnOperator<>(

              reduceFn,

              fullName,

              (Coder) windowedWorkItemCoder,

              mainTag,

              Collections.emptyList(),

              new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, outputCoder),

              windowingStrategy,

              new HashMap<>(), /* side-input mapping */

              Collections.emptyList(), /* side inputs */

              context.getPipelineOptions(),

              inputKvCoder.getKeyCoder(),

              keySelector);



      // our operator expects WindowedValue<KeyedWorkItem> while our input stream

      // is WindowedValue<SingletonKeyedWorkItem>, which is fine but Java doesn't like it ...

      @SuppressWarnings("unchecked")

      SingleOutputStreamOperator<WindowedValue<KV<K, Iterable<InputT>>>> outDataStream =

          keyedWorkItemStream

              .transform(fullName, outputTypeInfo, (OneInputStreamOperator) doFnOperator)

              .uid(fullName);



      context.setOutputDataStream(context.getOutput(transform), outDataStream);

    }

  }



  private static class CombinePerKeyTranslator<K, InputT, OutputT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>>> {



    @Override

    boolean canTranslate(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>> transform,

        GaiaStreamingTranslationContext context) {

      // if we have a merging window strategy and side inputs we cannot

      // translate as a proper combine. We have to group and then run the combine

      // over the final grouped values.

      PCollection<KV<K, InputT>> input = context.getInput(transform);



      @SuppressWarnings("unchecked")

      WindowingStrategy<?, BoundedWindow> windowingStrategy =

          (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();



      return windowingStrategy.getWindowFn().isNonMerging()

          || ((Combine.PerKey) transform).getSideInputs().isEmpty();

    }



    @Override

    public void translateNode(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KV<K, OutputT>>> transform,

        GaiaStreamingTranslationContext context) {

      String fullName = getCurrentTransformName(context);

      PCollection<KV<K, InputT>> input = context.getInput(transform);



      @SuppressWarnings("unchecked")

      WindowingStrategy<?, BoundedWindow> windowingStrategy =

          (WindowingStrategy<?, BoundedWindow>) input.getWindowingStrategy();



      KvCoder<K, InputT> inputKvCoder = (KvCoder<K, InputT>) input.getCoder();



      SingletonKeyedWorkItemCoder<K, InputT> workItemCoder =

          SingletonKeyedWorkItemCoder.of(

              inputKvCoder.getKeyCoder(),

              inputKvCoder.getValueCoder(),

              input.getWindowingStrategy().getWindowFn().windowCoder());



      DataStream<WindowedValue<KV<K, InputT>>> inputDataStream = context.getInputDataStream(input);



      WindowedValue.FullWindowedValueCoder<SingletonKeyedWorkItem<K, InputT>>

          windowedWorkItemCoder =

              WindowedValue.getFullCoder(

                  workItemCoder, input.getWindowingStrategy().getWindowFn().windowCoder());



      CoderTypeInformation<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemTypeInfo =

          new CoderTypeInformation<>(windowedWorkItemCoder);



      DataStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemStream =

          inputDataStream

              .flatMap(new ToKeyedWorkItem<>())

              .returns(workItemTypeInfo)

              .name("ToKeyedWorkItem");



      WorkItemKeySelector keySelector = new WorkItemKeySelector<>(inputKvCoder.getKeyCoder());

      KeyedStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>, ByteBuffer>

          keyedWorkItemStream = workItemStream.keyBy(keySelector);



      GlobalCombineFn<? super InputT, ?, OutputT> combineFn = ((Combine.PerKey) transform).getFn();

      SystemReduceFn<K, InputT, ?, OutputT, BoundedWindow> reduceFn =

          SystemReduceFn.combining(

              inputKvCoder.getKeyCoder(),

              AppliedCombineFn.withInputCoder(

                  combineFn, input.getPipeline().getCoderRegistry(), inputKvCoder));



      Coder<WindowedValue<KV<K, OutputT>>> outputCoder =

          context.getWindowedInputCoder(context.getOutput(transform));

      TypeInformation<WindowedValue<KV<K, OutputT>>> outputTypeInfo =

          context.getTypeInfo(context.getOutput(transform));



      List<PCollectionView<?>> sideInputs = ((Combine.PerKey) transform).getSideInputs();



      if (sideInputs.isEmpty()) {

        TupleTag<KV<K, OutputT>> mainTag = new TupleTag<>("main output");

        WindowDoFnOperator<K, InputT, OutputT> doFnOperator =

            new WindowDoFnOperator<>(

                reduceFn,

                fullName,

                (Coder) windowedWorkItemCoder,

                mainTag,

                Collections.emptyList(),

                new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, outputCoder),

                windowingStrategy,

                new HashMap<>(), /* side-input mapping */

                Collections.emptyList(), /* side inputs */

                context.getPipelineOptions(),

                inputKvCoder.getKeyCoder(),

                keySelector);



        // our operator expects WindowedValue<KeyedWorkItem> while our input stream

        // is WindowedValue<SingletonKeyedWorkItem>, which is fine but Java doesn't like it ...

        @SuppressWarnings("unchecked")

        SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> outDataStream =

            keyedWorkItemStream

                .transform(fullName, outputTypeInfo, (OneInputStreamOperator) doFnOperator)

                .uid(fullName);

        context.setOutputDataStream(context.getOutput(transform), outDataStream);

      } else {

        Tuple2<Map<Integer, PCollectionView<?>>, DataStream<RawUnionValue>> transformSideInputs =

            transformSideInputs(sideInputs, context);



        TupleTag<KV<K, OutputT>> mainTag = new TupleTag<>("main output");

        WindowDoFnOperator<K, InputT, OutputT> doFnOperator =

            new WindowDoFnOperator<>(

                reduceFn,

                fullName,

                (Coder) windowedWorkItemCoder,

                mainTag,

                Collections.emptyList(),

                new DoFnOperator.MultiOutputOutputManagerFactory<>(mainTag, outputCoder),

                windowingStrategy,

                transformSideInputs.f0,

                sideInputs,

                context.getPipelineOptions(),

                inputKvCoder.getKeyCoder(),

                keySelector);



        // we have to manually contruct the two-input transform because we're not

        // allowed to have only one input keyed, normally.



        TwoInputTransformation<

                WindowedValue<SingletonKeyedWorkItem<K, InputT>>,

                RawUnionValue,

                WindowedValue<KV<K, OutputT>>>

            rawGaiaTransform =

                new TwoInputTransformation<>(

                    keyedWorkItemStream.getTransformation(),

                    transformSideInputs.f1.broadcast().getTransformation(),

                    transform.getName(),

                    (TwoInputStreamOperator) doFnOperator,

                    outputTypeInfo,

                    keyedWorkItemStream.getParallelism());



        rawGaiaTransform.setStateKeyType(keyedWorkItemStream.getKeyType());

        rawGaiaTransform.setStateKeySelectors(keyedWorkItemStream.getKeySelector(), null);



        @SuppressWarnings({"unchecked", "rawtypes"})

        SingleOutputStreamOperator<WindowedValue<KV<K, OutputT>>> outDataStream =

            new SingleOutputStreamOperator(

                keyedWorkItemStream.getExecutionEnvironment(),

                rawGaiaTransform) {}; // we have to cheat around the ctor being protected



        keyedWorkItemStream.getExecutionEnvironment().addOperator(rawGaiaTransform);



        context.setOutputDataStream(context.getOutput(transform), outDataStream);

      }

    }

  }



  private static class GBKIntoKeyedWorkItemsTranslator<K, InputT>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<KV<K, InputT>>, PCollection<KeyedWorkItem<K, InputT>>>> {



    @Override

    boolean canTranslate(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KeyedWorkItem<K, InputT>>> transform,

        GaiaStreamingTranslationContext context) {

      return true;

    }



    @Override

    public void translateNode(

        PTransform<PCollection<KV<K, InputT>>, PCollection<KeyedWorkItem<K, InputT>>> transform,

        GaiaStreamingTranslationContext context) {



      PCollection<KV<K, InputT>> input = context.getInput(transform);



      KvCoder<K, InputT> inputKvCoder = (KvCoder<K, InputT>) input.getCoder();



      SingletonKeyedWorkItemCoder<K, InputT> workItemCoder =

          SingletonKeyedWorkItemCoder.of(

              inputKvCoder.getKeyCoder(),

              inputKvCoder.getValueCoder(),

              input.getWindowingStrategy().getWindowFn().windowCoder());



      WindowedValue.ValueOnlyWindowedValueCoder<SingletonKeyedWorkItem<K, InputT>>

          windowedWorkItemCoder = WindowedValue.getValueOnlyCoder(workItemCoder);



      CoderTypeInformation<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemTypeInfo =

          new CoderTypeInformation<>(windowedWorkItemCoder);



      DataStream<WindowedValue<KV<K, InputT>>> inputDataStream = context.getInputDataStream(input);



      DataStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> workItemStream =

          inputDataStream

              .flatMap(new ToKeyedWorkItemInGlobalWindow<>())

              .returns(workItemTypeInfo)

              .name("ToKeyedWorkItem");



      KeyedStream<WindowedValue<SingletonKeyedWorkItem<K, InputT>>, ByteBuffer>

          keyedWorkItemStream =

              workItemStream.keyBy(new WorkItemKeySelector<>(inputKvCoder.getKeyCoder()));



      context.setOutputDataStream(context.getOutput(transform), keyedWorkItemStream);

    }

  }



  private static class ToKeyedWorkItemInGlobalWindow<K, InputT>

      extends RichFlatMapFunction<

          WindowedValue<KV<K, InputT>>, WindowedValue<SingletonKeyedWorkItem<K, InputT>>> {



    @Override

    public void flatMap(

        WindowedValue<KV<K, InputT>> inWithMultipleWindows,

        Collector<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> out)

        throws Exception {



      // we need to wrap each one work item per window for now

      // since otherwise the PushbackSideInputRunner will not correctly

      // determine whether side inputs are ready

      //

      // this is tracked as https://issues.apache.org/jira/browse/BEAM-1850

      for (WindowedValue<KV<K, InputT>> in : inWithMultipleWindows.explodeWindows()) {

        SingletonKeyedWorkItem<K, InputT> workItem =

            new SingletonKeyedWorkItem<>(

                in.getValue().getKey(), in.withValue(in.getValue().getValue()));



        out.collect(WindowedValue.valueInGlobalWindow(workItem));

      }

    }

  }



  private static class FlattenPCollectionTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<

          PTransform<PCollection<T>, PCollection<T>>> {



    @Override

    public void translateNode(

        PTransform<PCollection<T>, PCollection<T>> transform,

        GaiaStreamingTranslationContext context) {

      Map<TupleTag<?>, PValue> allInputs = context.getInputs(transform);



      if (allInputs.isEmpty()) {



        // create an empty dummy source to satisfy downstream operations

        // we cannot create an empty source in Gaia, therefore we have to

        // add the flatMap that simply never forwards the single element

        DataStreamSource<String> dummySource =

            context.getExecutionEnvironment().fromElements("dummy");



        DataStream<WindowedValue<T>> result =

            dummySource

                .<WindowedValue<T>>flatMap(

                    (s, collector) -> {

                      // never return anything

                    })

                .returns(

                    new CoderTypeInformation<>(

                        WindowedValue.getFullCoder(

                            (Coder<T>) VoidCoder.of(), GlobalWindow.Coder.INSTANCE)));

        context.setOutputDataStream(context.getOutput(transform), result);



      } else {

        DataStream<T> result = null;



        // Determine DataStreams that we use as input several times. For those, we need to uniquify

        // input streams because Gaia seems to swallow watermarks when we have a union of one and

        // the same stream.

        Map<DataStream<T>, Integer> duplicates = new HashMap<>();

        for (PValue input : allInputs.values()) {

          DataStream<T> current = context.getInputDataStream(input);

          Integer oldValue = duplicates.put(current, 1);

          if (oldValue != null) {

            duplicates.put(current, oldValue + 1);

          }

        }



        for (PValue input : allInputs.values()) {

          DataStream<T> current = context.getInputDataStream(input);



          final Integer timesRequired = duplicates.get(current);

          if (timesRequired > 1) {

            current =

                current.flatMap(

                    new FlatMapFunction<T, T>() {

                      private static final long serialVersionUID = 1L;



                      @Override

                      public void flatMap(T t, Collector<T> collector) throws Exception {

                        collector.collect(t);

                      }

                    });

          }

          result = (result == null) ? current : result.union(current);

        }



        context.setOutputDataStream(context.getOutput(transform), result);

      }

    }

  }



  static class ToKeyedWorkItem<K, InputT>

      extends RichFlatMapFunction<

          WindowedValue<KV<K, InputT>>, WindowedValue<SingletonKeyedWorkItem<K, InputT>>> {



    @Override

    public void flatMap(

        WindowedValue<KV<K, InputT>> inWithMultipleWindows,

        Collector<WindowedValue<SingletonKeyedWorkItem<K, InputT>>> out)

        throws Exception {



      // we need to wrap each one work item per window for now

      // since otherwise the PushbackSideInputRunner will not correctly

      // determine whether side inputs are ready

      //

      // this is tracked as https://issues.apache.org/jira/browse/BEAM-1850

      for (WindowedValue<KV<K, InputT>> in : inWithMultipleWindows.explodeWindows()) {

        SingletonKeyedWorkItem<K, InputT> workItem =

            new SingletonKeyedWorkItem<>(

                in.getValue().getKey(), in.withValue(in.getValue().getValue()));



        out.collect(in.withValue(workItem));

      }

    }

  }



  /**

   * A translator just to vend the URN. This will need to be moved to runners-core-construction-java

   * once SDF is reorganized appropriately.

   */

  private static class SplittableParDoProcessElementsTranslator

      extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<

          SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?>> {



    private SplittableParDoProcessElementsTranslator() {}



    @Override

    public String getUrn(SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?> transform) {

      return SPLITTABLE_PROCESS_URN;

    }

  }



  /** Registers classes specialized to the Gaia runner. */

  @AutoService(TransformPayloadTranslatorRegistrar.class)

  public static class GaiaTransformsRegistrar implements TransformPayloadTranslatorRegistrar {

    @Override

    public Map<

            ? extends Class<? extends PTransform>,

            ? extends PTransformTranslation.TransformPayloadTranslator>

        getTransformPayloadTranslators() {

      return ImmutableMap

          .<Class<? extends PTransform>, PTransformTranslation.TransformPayloadTranslator>builder()

          .put(

              CreateStreamingGaiaView.CreateGaiaPCollectionView.class,

              new CreateStreamingGaiaViewPayloadTranslator())

          .put(

              SplittableParDoViaKeyedWorkItems.ProcessElements.class,

              new SplittableParDoProcessElementsTranslator())

          .put(

              SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems.class,

              new SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator())

          .build();

    }

  }



  /**

   * A translator just to vend the URN. This will need to be moved to runners-core-construction-java

   * once SDF is reorganized appropriately.

   */

  private static class SplittableParDoProcessElementsPayloadTranslator

      extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<

          SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?>> {



    private SplittableParDoProcessElementsPayloadTranslator() {}



    @Override

    public String getUrn(SplittableParDoViaKeyedWorkItems.ProcessElements<?, ?, ?, ?> transform) {

      return SPLITTABLE_PROCESS_URN;

    }

  }



  /**

   * A translator just to vend the URN. This will need to be moved to runners-core-construction-java

   * once SDF is reorganized appropriately.

   */

  private static class SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator

      extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<

          SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<?, ?>> {



    private SplittableParDoGbkIntoKeyedWorkItemsPayloadTranslator() {}



    @Override

    public String getUrn(SplittableParDoViaKeyedWorkItems.GBKIntoKeyedWorkItems<?, ?> transform) {

      return SplittableParDo.SPLITTABLE_GBKIKWI_URN;

    }

  }



  /** A translator just to vend the URN. */

  private static class CreateStreamingGaiaViewPayloadTranslator

      extends PTransformTranslation.TransformPayloadTranslator.NotSerializable<

          CreateStreamingGaiaView.CreateGaiaPCollectionView<?, ?>> {



    private CreateStreamingGaiaViewPayloadTranslator() {}



    @Override

    public String getUrn(CreateStreamingGaiaView.CreateGaiaPCollectionView<?, ?> transform) {

      return CreateStreamingGaiaView.CREATE_STREAMING_GAIA_VIEW_URN;

    }

  }



  /** A translator to support {@link TestStream} with Gaia. */

  private static class TestStreamTranslator<T>

      extends GaiaStreamingPipelineTranslator.StreamTransformTranslator<TestStream<T>> {



    @Override

    void translateNode(TestStream<T> testStream, GaiaStreamingTranslationContext context) {

      Coder<T> valueCoder = testStream.getValueCoder();



      // Coder for the Elements in the TestStream

      TestStream.TestStreamCoder<T> testStreamCoder = TestStream.TestStreamCoder.of(valueCoder);

      final byte[] payload;

      try {

        payload = CoderUtils.encodeToByteArray(testStreamCoder, testStream);

      } catch (CoderException e) {

        throw new RuntimeException("Could not encode TestStream.", e);

      }



      SerializableFunction<byte[], TestStream<T>> testStreamDecoder =

          bytes -> {

            try {

              return CoderUtils.decodeFromByteArray(

                  TestStream.TestStreamCoder.of(valueCoder), bytes);

            } catch (CoderException e) {

              throw new RuntimeException("Can't decode TestStream payload.", e);

            }

          };



      WindowedValue.FullWindowedValueCoder<T> elementCoder =

          WindowedValue.getFullCoder(valueCoder, GlobalWindow.Coder.INSTANCE);



      DataStreamSource<WindowedValue<T>> source =

          context

              .getExecutionEnvironment()

              .addSource(

                  new TestStreamSource<>(testStreamDecoder, payload),

                  new CoderTypeInformation<>(elementCoder));



      context.setOutputDataStream(context.getOutput(testStream), source);

    }

  }



  /**

   * Wrapper for {@link UnboundedSourceWrapper}, which simplifies output type, namely, removes

   * {@link ValueWithRecordId}.

   */

  static class UnboundedSourceWrapperNoValueWithRecordId<

          OutputT, CheckpointMarkT extends UnboundedSource.CheckpointMark>

      extends RichParallelSourceFunction<WindowedValue<OutputT>>

      implements ProcessingTimeCallback,

          StoppableFunction,

          CheckpointListener,

          CheckpointedFunction {



    private final UnboundedSourceWrapper<OutputT, CheckpointMarkT> unboundedSourceWrapper;



    @VisibleForTesting

	UnboundedSourceWrapper<OutputT, CheckpointMarkT> getUnderlyingSource() {

      return unboundedSourceWrapper;

    }



    UnboundedSourceWrapperNoValueWithRecordId(

        UnboundedSourceWrapper<OutputT, CheckpointMarkT> unboundedSourceWrapper) {

      this.unboundedSourceWrapper = unboundedSourceWrapper;

    }



    @Override

    public void open(Configuration parameters) throws Exception {

      unboundedSourceWrapper.setRuntimeContext(getRuntimeContext());

      unboundedSourceWrapper.open(parameters);

    }



    @Override

    public void run(SourceContext<WindowedValue<OutputT>> ctx) throws Exception {

      unboundedSourceWrapper.run(new SourceContextWrapper(ctx));

    }



    @Override

    public void initializeState(FunctionInitializationContext context) throws Exception {

      unboundedSourceWrapper.initializeState(context);

    }



    @Override

    public void snapshotState(FunctionSnapshotContext context) throws Exception {

      unboundedSourceWrapper.snapshotState(context);

    }



    @Override

    public void notifyCheckpointComplete(long checkpointId) throws Exception {

      unboundedSourceWrapper.notifyCheckpointComplete(checkpointId);

    }



    @Override

    public void stop() {

      unboundedSourceWrapper.stop();

    }



    @Override

    public void cancel() {

      unboundedSourceWrapper.cancel();

    }



    @Override

    public void onProcessingTime(long timestamp) throws Exception {

      unboundedSourceWrapper.onProcessingTime(timestamp);

    }



    private final class SourceContextWrapper

        implements SourceContext<WindowedValue<ValueWithRecordId<OutputT>>> {



      private final SourceContext<WindowedValue<OutputT>> ctx;



      private SourceContextWrapper(SourceContext<WindowedValue<OutputT>> ctx) {

        this.ctx = ctx;

      }



      @Override

      public void collect(WindowedValue<ValueWithRecordId<OutputT>> element) {

        OutputT originalValue = element.getValue().getValue();

        WindowedValue<OutputT> output =

            WindowedValue.of(

                originalValue, element.getTimestamp(), element.getWindows(), element.getPane());

        ctx.collect(output);

      }



      @Override

      public void collectWithTimestamp(

          WindowedValue<ValueWithRecordId<OutputT>> element, long timestamp) {

        OutputT originalValue = element.getValue().getValue();

        WindowedValue<OutputT> output =

            WindowedValue.of(

                originalValue, element.getTimestamp(), element.getWindows(), element.getPane());

        ctx.collectWithTimestamp(output, timestamp);

      }



      @Override

      public void emitWatermark(Watermark mark) {

        ctx.emitWatermark(mark);

      }



      @Override

      public void markAsTemporarilyIdle() {

        ctx.markAsTemporarilyIdle();

      }



      @Override

      public Object getCheckpointLock() {

        return ctx.getCheckpointLock();

      }



      @Override

      public void close() {

        ctx.close();

      }

    }

  }

}